//
//  BinaryHeap.swift
//  DataStructureAndAlgorithm
//
//  Created by icemelon on 2021/9/15.
//

import Foundation

enum BinaryHeapType {
    case Min //最小堆
    case Max //最大堆
}
/** 索引index 的规律
 i = 0 root  父节点 floor((i - 1)/2)向下取整
 */
class BinaryHeap<T: Comparable> {
    lazy private var heap = Array<T>() // 用数组来存储数据
    private var type: BinaryHeapType
    
    public init(type: BinaryHeapType = .Max) {
        self.type = type
    }
    
    var size: Int {
        get { return heap.count }
    }
    
    var isEmpty: Bool {
        get { return heap.isEmpty }
    }
    
    func clear() {
        heap.removeAll()
    }
    
    func get() -> T? {
        if !isEmpty {
            return heap.first!
        }
        return nil 
    }
    
    func log() {
        print("\n heap: isEmpty\(isEmpty) size:\(size) get:\(get()!)")
        for t in heap {
            print(t,separator: " ",terminator: " ")
        }
    }
}

extension BinaryHeap {
    func add(element: T) {
        heap.append(element)
        siftUp(index: size - 1)//添加之后将添加的元素上滤
    }
    
    @discardableResult func remove() -> T? {
        if !isEmpty {
            //用最后一个几点覆盖根节点
            let first = heap.first!
            heap[0] = heap.last!
            //删除最后一个几点
            heap.removeLast()
            //下滤
            shifDown(index: 0)
            return first
        }
        return nil
    }
    
    @discardableResult func replace(element: T) -> T? {//替换根节点 然后下滤
        if !isEmpty {
            let first = heap.first!
            heap[0] = element
            if size > 1 { shifDown(index: 0) }
            return first
        }
        else {
            heap[0] = element
        }
        return nil
    }
    
}

extension BinaryHeap {
    private func siftUp(index: Int) {
        let addE = heap[index]   //toUpElement 保存一下需要上滤的元素
        var i = index
        while i > 0 {
            let parentIndex = getParentIndex(i)
            let parentE = heap[parentIndex]
            let result = compare(e1: addE, e2: parentE)
            if result < 0 { break } //无论是最大堆还是最小堆
            //parent的值覆盖需要上滤的位置
            heap[i] = parentE
            i = parentIndex
        }
        heap[i] = addE
    }
    
    private func shifDown(index: Int) {
        let rootElement = heap[index] //保存下滤的元素值
        var i = index
        let half = size / 2
        // 第一个叶子节点的索引 == 非叶子节点的数量
        // index < 第一个叶子节点的索引
        while index < half {// 必须保证index位置是非叶子节点
            guard let compareIndex = compareChildIndex(i: i) else {
                break
            }
            let result = compare(e1: rootElement, e2: heap[compareIndex])
            if result > 0 { return }
            heap[i] = heap[compareIndex]
            i = compareIndex
        }
        heap[i] = rootElement
    }
    
    private func compare(e1: T, e2: T) -> Int { //e1为添加的元素 e2是参与比较的节点 返回值为正数则表示需要替换 可以将函数返回 bool
        if self.type == .Max {
            return e1 > e2 ? 1 : -1
        }
        else {
            return e1 < e2 ? 1 : -1
        }
    }
    
    private func getParentIndex(_ current: Int) -> Int {
        return (current - 1) / 2  //floor(<#T##Double#>) 向下取整
    }
    
    private func leftChildIndex(i: Int) -> Int {
        return i * 2 + 1
    }
    
    private func leftChild(i: Int) -> T? {
        let index = leftChildIndex(i: i)
        if index < size {
            return heap[index]
        }
        else {
            return nil
        }
    }
    
    private func rightChildIndex(i: Int) -> Int {
        return i * 2 + 2
    }
    
    private func rightChild(i: Int) -> T? {
        let index = rightChildIndex(i: i)
        if index < size {
            return heap[index]
        }
        else {
            return nil
        }
    }
    
    private func compareChildIndex(i: Int) -> Int? { //获取应该参与比较的子节点的index
        var index = leftChildIndex(i: i)
        guard let child = leftChild(i: i) else {
            return nil
        }
        let right = rightChild(i: i)
        if right != nil {
            if type == .Max { //区分最大堆还是最小堆
                if right! > child {
                    index = rightChildIndex(i: i)
                }
            }
            else {
                if right! < child {
                    index = rightChildIndex(i: i)
                }
            }
            
        }
        return index
    }
}
 
extension BinaryHeap {
    class func test() {
        let b = BinaryHeap<Int>(type: .Max)
        let array = random()
        
        for num in array {
            b.add(element: num)
        }
        b.log()
        b.remove()
        b.log()
        b.remove()
        b.log()

        b.replace(element: 999)
        b.log()
        b.replace(element: 666)
        b.log()
        b.replace(element: 222)
        b.log()
        b.replace(element: 111)
        b.log()
    }
    class func TopK() {
        let array = random(count: 10, max: 100)
        print(array, separator: " ", terminator: " ")
        let bh = topK(data: array, k: 2)
        bh.log()
    }
}
